#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
模型训练器模块
"""

import os
import json
import time
import threading
import numpy as np
from pathlib import Path

class ModelTrainer(threading.Thread):
    """模型训练线程类"""
    
    def __init__(self, dataset_path, model_type, params, logger):
        """
        初始化训练器
        
        Args:
            dataset_path (str): 数据集路径
            model_type (str): 模型类型（如 'classification', 'detection'）
            params (dict): 训练参数
            logger (logging.Logger): 日志器
        """
        super().__init__()
        self.dataset_path = Path(dataset_path)
        self.model_type = model_type
        self.params = params
        self.logger = logger
        
        # 训练状态
        self.running = False
        self.should_stop = False
        
        # 回调函数
        self.on_progress = None  # 进度回调，参数：(epoch, metrics)
        self.on_completed = None  # 完成回调，参数：(success, metrics)
        
        # 训练结果
        self.model = None
    
    def run(self):
        """线程运行方法"""
        self.running = True
        self.should_stop = False
        success = False
        final_metrics = {}
        
        try:
            # 根据模型类型选择不同的训练方法
            if self.model_type == "classification":
                self.model, final_metrics = self._train_classification_model()
                success = True
            elif self.model_type == "detection":
                self.model, final_metrics = self._train_detection_model()
                success = True
            else:
                self.logger.error(f"不支持的模型类型: {self.model_type}")
        except Exception as e:
            self.logger.error(f"训练失败: {e}")
            final_metrics = {"error": str(e)}
        
        self.running = False
        
        # 调用完成回调
        if self.on_completed:
            self.on_completed(success, final_metrics)
    
    def stop(self):
        """停止训练"""
        self.should_stop = True
    
    def _load_dataset(self):
        """
        加载数据集
        
        Returns:
            tuple: (images, labels, classes)
        """
        self.logger.info("加载数据集...")
        
        # 读取数据集元数据
        metadata_file = self.dataset_path / "dataset_meta.json"
        with open(metadata_file, "r", encoding="utf-8") as f:
            metadata = json.load(f)
        
        # 获取图像和标注路径
        images_dir = self.dataset_path / "images"
        annotations_dir = self.dataset_path / "annotations"
        
        # 获取所有图像文件
        image_files = []
        for f in images_dir.glob("*"):
            if f.is_file() and f.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp']:
                image_files.append(f)
        
        # 对于分类模型，我们需要图像和类别标签
        images = []
        labels = []
        classes = []
        
        # 获取所有类别
        for label in metadata.get("labels", []):
            classes.append(label["name"])
        
        # 读取图像和标注
        for img_file in image_files:
            # 检查是否有相应的标注
            anno_file = annotations_dir / f"{img_file.stem}.json"
            if anno_file.exists():
                with open(anno_file, "r", encoding="utf-8") as f:
                    annotations = json.load(f)
                
                # 对于分类，我们只使用第一个标注的类别
                if annotations:
                    label = annotations[0].get("label", "")
                    if label in classes:
                        import cv2
                        img = cv2.imread(str(img_file))
                        if img is not None:
                            images.append(img)
                            labels.append(classes.index(label))
        
        self.logger.info(f"已加载 {len(images)} 张图像，{len(classes)} 个类别")
        return np.array(images), np.array(labels), classes
    
    def _train_classification_model(self):
        """
        训练分类模型
        
        Returns:
            tuple: (model, metrics)
        """
        self.logger.info("开始训练分类模型...")
        
        # 加载数据集
        images, labels, classes = self._load_dataset()
        
        # 如果没有足够的数据，返回错误
        if len(images) < 10 or len(classes) < 2:
            self.logger.error("数据不足，无法训练")
            return None, {"error": "数据不足，无法训练"}
        
        # 获取训练参数
        batch_size = self.params.get("batch_size", 32)
        learning_rate = self.params.get("learning_rate", 0.001)
        epochs = self.params.get("epochs", 50)
        val_split = self.params.get("val_split", 0.2)
        algorithm = self.params.get("algorithm", "efficientnet")  # 默认使用EfficientNet
        
        # 记录选择的算法
        self.logger.info(f"选择的训练算法: {algorithm}")
        
        # 根据不同算法类型初始化不同的模型
        # 实际应用中应该导入对应的库并实例化模型
        model_info = {"type": "classification", "algorithm": algorithm, "classes": classes}
        
        # 输出算法特点
        if algorithm == "efficientnet":
            self.logger.info("使用EfficientNet - 平衡的精度与速度")
        elif algorithm == "mobilenetv3":
            self.logger.info("使用MobileNetV3 - 适合移动设备的轻量级模型")
        elif algorithm == "resnet":
            self.logger.info("使用ResNet - 高精度残差网络")
        elif algorithm == "vit":
            self.logger.info("使用Vision Transformer - 基于Transformer架构的视觉模型")
        
        # 在这里可以添加实际的训练代码
        # 这里只是一个示例，实际应该使用PyTorch或TensorFlow进行训练
        
        # 模拟训练过程
        model = {"type": "classification", "algorithm": algorithm, "classes": classes}
        
        # 训练指标
        train_losses = []
        val_losses = []
        train_accs = []
        val_accs = []
        
        for epoch in range(epochs):
            if self.should_stop:
                self.logger.info("训练被用户停止")
                break
            
            # 模拟训练一个epoch
            time.sleep(0.5)  # 模拟训练时间
            
            # 计算模拟指标
            train_loss = 1.0 / (epoch + 1)
            val_loss = 1.2 / (epoch + 1)
            train_acc = min(95.0, 50.0 + epoch * 1.0)
            val_acc = min(90.0, 45.0 + epoch * 1.0)
            
            train_losses.append(train_loss)
            val_losses.append(val_loss)
            train_accs.append(train_acc)
            val_accs.append(val_acc)
            
            # 调用进度回调
            if self.on_progress:
                # 创建度量指标的副本
                metrics = {
                    "train_loss": float(train_loss),
                    "val_loss": float(val_loss),
                    "train_acc": float(train_acc),
                    "val_acc": float(val_acc),
                    "learning_rate": float(learning_rate)
                }
                
                # 使用整数epoch而不是引用对象
                self.on_progress(int(epoch), metrics)
        
        # 创建一个模拟的混淆矩阵
        n_classes = len(classes)
        confusion_matrix = np.zeros((n_classes, n_classes), dtype=np.int32)
        for i in range(n_classes):
            for j in range(n_classes):
                if i == j:
                    confusion_matrix[i, j] = np.random.randint(5, 20)
                else:
                    confusion_matrix[i, j] = np.random.randint(0, 5)
        
        # 最终指标
        final_metrics = {
            "train_loss": float(train_losses[-1]) if train_losses else 0.0,
            "val_loss": float(val_losses[-1]) if val_losses else 0.0,
            "train_acc": float(train_accs[-1]) if train_accs else 0.0,
            "val_acc": float(val_accs[-1]) if val_accs else 0.0,
            "confusion_matrix": confusion_matrix.copy(),
            "labels": classes.copy() if isinstance(classes, list) else classes.tolist()
        }
        
        self.logger.info("分类模型训练完成")
        return model, final_metrics
    
    def _train_detection_model(self):
        """
        训练目标检测模型
        
        Returns:
            tuple: (model, metrics)
        """
        self.logger.info("开始训练目标检测模型...")
        
        # 加载数据集
        images, labels, classes = self._load_dataset()
        
        # 如果没有足够的数据，返回错误
        if len(images) < 10 or len(classes) < 2:
            self.logger.error("数据不足，无法训练")
            return None, {"error": "数据不足，无法训练"}
        
        # 获取训练参数
        batch_size = self.params.get("batch_size", 32)
        learning_rate = self.params.get("learning_rate", 0.001)
        epochs = self.params.get("epochs", 50)
        val_split = self.params.get("val_split", 0.2)
        algorithm = self.params.get("algorithm", "yolov8")  # 默认使用YOLOv8
        
        # 记录选择的算法
        self.logger.info(f"选择的训练算法: {algorithm}")
        
        # 根据不同算法类型初始化不同的模型
        if algorithm == "yolov8":
            self.logger.info("使用YOLOv8 - 实时目标检测算法")
        elif algorithm == "fasterrcnn":
            self.logger.info("使用Faster R-CNN - 高精度两阶段检测算法")
        elif algorithm == "nanodet":
            self.logger.info("使用NanoDet - 超轻量级检测算法，适合边缘设备")
        
        # 模拟训练过程
        model = {"type": "detection", "algorithm": algorithm, "classes": classes}
        
        # 训练指标
        train_losses = []
        val_losses = []
        
        for epoch in range(epochs):
            if self.should_stop:
                self.logger.info("训练被用户停止")
                break
            
            # 模拟训练一个epoch
            time.sleep(0.5)  # 模拟训练时间
            
            # 计算模拟指标
            train_loss = 1.0 / (epoch + 1)
            val_loss = 1.2 / (epoch + 1)
            
            train_losses.append(train_loss)
            val_losses.append(val_loss)
            
            # 调用进度回调
            if self.on_progress:
                metrics = {
                    "train_loss": float(train_loss),
                    "val_loss": float(val_loss),
                    "learning_rate": float(learning_rate)
                }
                self.on_progress(int(epoch), metrics)
        
        # 最终指标
        final_metrics = {
            "train_loss": float(train_losses[-1]) if train_losses else 0.0,
            "val_loss": float(val_losses[-1]) if val_losses else 0.0
        }
        
        self.logger.info("目标检测模型训练完成")
        return model, final_metrics
    
    def export_model(self, export_path, format="onnx"):
        """
        导出模型
        
        Args:
            export_path (str): 导出路径
            format (str): 导出格式，如 'onnx', 'pt', 'h5'
            
        Returns:
            bool: 导出是否成功
        """
        if not self.model:
            self.logger.error("没有训练好的模型可导出")
            return False
        
        try:
            self.logger.info(f"导出模型到 {export_path}，格式: {format}")
            
            # 模拟导出过程
            time.sleep(1)
            
            # 创建一个空的导出文件
            with open(export_path, "w") as f:
                json.dump({"type": "model", "format": format}, f)
            
            self.logger.info("模型导出成功")
            return True
        except Exception as e:
            self.logger.error(f"导出模型失败: {e}")
            return False 